from numpy.core.function_base import linspace
import scipy.io as scio
import numpy as np 
import matplotlib.pyplot as plt
import matplotlib.image as img
from PIL import Image
from matplotlib.backends.backend_agg import FigureCanvasAgg
from PIL import ImageDraw
import os
import concurrent.futures
import time
import traceback
def rgb2gray(rgb):
        return np.dot(rgb[...,:3], [0.299, 0.587, 0.114])
def read_wall(path, k=1, dr=0, dz=0):
    mask = img.imread(path)
    #mask = mask.resize((65,65))
    mask = np.array(mask)
    mask = rgb2gray(mask)
    [z, r] = np.shape(mask)
    Z = np.linspace(0.96, -0.96, z)
    R = np.linspace(1.2, 2.28, r)
    RX = []
    ZY = []
    for rr in range(r):
        for zz in range(z):
            if mask[zz, rr] == 0:
                RX.append(R[rr])
                ZY.append(Z[zz])
    RX = np.array(RX)
    ZY = np.array(ZY)
    RX = k*RX + dr
    ZY = k*ZY + dz
    plt.plot(RX,ZY,',',color = 'w')

    #mask[mask!=0] = 1
    #plt.imshow(mask, cmap='Greys_r')
    #plt.contour(R, Z, mask)
    return mask
def show_pic(path):
    pic = img.imread(path)
    pic = np.array(pic)
    pic = np.fliplr(pic)
    pic = np.flipud(pic)
    plt.imshow(pic)
def image_cut(path,vars,savepath):
    [l,t,r,b] = vars
    # 图片路径，相对路径
    #image_path = 'H:\\挑战杯\\mak-HL2A\\CCD\\21400\\FCCD\\21400\\jpg\\00015.jpg'
    # 读取图片
    image = Image.open(path)
    image= image.transpose(Image.FLIP_LEFT_RIGHT)
    image= image.transpose(Image.FLIP_TOP_BOTTOM)
    draw = ImageDraw.Draw(image)
    #draw.rectangle(((150,30),(600,480)), outline='white')
    image = image.crop((l,t,r,b))
    # 输出图片维度
    #print("image_shape: ", image.size)
    image= image.transpose(Image.FLIP_TOP_BOTTOM)
    # 显示图片
    #image.show()
    image.save(savepath)
    return image.size
def reshape_pic(path,var):
    var = np.array(var, dtype=np.int)
    
    pic = Image.open(path)
    pic = pic.resize(var)
    pic.save(path)
def gen_png(shoot, frame, vars):
    [k, dr, dz] = vars
    pic_path =  'jpg\\'+shoot+'\\jpg\\'+frame+'.jpg'
    label_path = 'cut\\efit\\'+shoot+frame+'.png'
    pic_savepath = 'cut\\pic\\'+shoot+frame+'.png'
    try:
    # 不能确定正确执行的代码
        data = scio.loadmat('efit\\'+shoot+'\\'+frame+'.mat')
    except:
        print('no efit at shoot '+ shoot + ' frame ' + frame)
        return
    efit = data['ansl']
    psibry = efit['psi'][0][0][0][0]
    psimap = efit['psimap'][0][0]
    rmap = efit['rmap'][0][0]
    zmap = efit['zmap'][0][0]
    psimax = np.max(psimap)
    R, Z = np.meshgrid(rmap, zmap)
    R = k*R + dr
    Z = k*Z + dz
    #plt.figure()
    #plt.axis('off')
    plt.contourf(R, Z, psimap, [-100, psibry, psimax], colors = [(0,0.0,0.0),(0.5,0.0,0.0)] )
    plt.axis('image')
    plt.axis('off')
    xmin, xmax = plt.xlim()
    ymin, ymax = plt.ylim()
    plt.ylim([0,ymax])
    #print('efit_shape: ' + str([xmax-xmin, ymax-0]))
    plt.savefig(label_path, bbox_inches='tight', pad_inches=0)
    size = image_cut(pic_path, [xmin, 0, xmax, ymax], pic_savepath)
    reshape_pic(label_path,size)
    #plt.show()

def viewmak(shoot, frame, vars):
    [k, dr, dz] = vars
    
    pic_path = 'jpg\\'+shoot+'\\jpg\\'+frame+'.jpg'
    try:
    # 不能确定正确执行的代码
        data = scio.loadmat('efit\\'+shoot+'\\'+frame+'.mat')
    except:
        print('no efit at shoot '+ shoot + ' frame ' + frame)
        
        try:
            show_pic(pic_path)
        except:
            print('and no pic')
            return 1
        read_wall('wall.png',k,dr,dz)
        plt.axis('equal')
        ax = plt.gca()
        ax.invert_yaxis()
        ax.yaxis.set_ticks_position('left')
        plt.show()
        return
    
    
    efit = data['ansl']

    psibry = efit['psi'][0][0][0][0]
    psimap = efit['psimap'][0][0]
    rmap = efit['rmap'][0][0]
    zmap = efit['zmap'][0][0]
    psimax = np.max(psimap)
    R, Z = np.meshgrid(rmap, zmap)
    #print('psibry: ' + str(psibry))
    #print('psimap: ' + str(psimap))
    #print('psimax: ' + str(psimax))
    #print('zmap: ' + str(zmap))
    #print('rmap: ' + str(rmap))
    #print('zmap shape: ' + str(np.shape(zmap)))
    #print('rmap shape: ' + str(np.shape(rmap)))
    plt.axis('equal')

    
    
    show_pic(pic_path)
    read_wall('wall.png',k,dr,dz)
    #data = scio.loadmat('wall.fig')
    #print(data)
    R = k*R + dr
    Z = k*Z + dz
    #plt.text(0.5*k+dr,0.75*k+dz, 'psibry: '+str(psibry), color = 'y')
    plt.contour(R, Z, psimap, [psibry], colors = [(0.5,0.0,0.0)] )
    ax = plt.gca()
    ax.invert_yaxis()
    ax.yaxis.set_ticks_position('left')
    #plt.axis('off')




    plt.show()
def view_label(shoot,frame,save=0):
    plt.cla()
    label_path = 'cut\\efit\\'+shoot+frame+'.png'
    pic_path = 'cut\\pic\\'+shoot+frame+'.png'
    pic_savepath = 'cut\\judge\\'+shoot+frame+'.png'
    pic =  np.array(img.imread(pic_path))
    plt.imshow(pic)
    label =  np.array(img.imread(label_path))
    plt.imshow(label, alpha=0.5)
    plt.axis('image')
    plt.axis('off')
    
    if save == 1:
        plt.savefig(pic_savepath, bbox_inches='tight', pad_inches=0)
    else:
        plt.show()
def get_efit(shoot,frame):

    frame = int(frame)
    framestd = '%05d'%frame
    dr = -150
    k = 280
    dz = 0.96*k-70
    vars = [k, dr, dz]
    '''
    plt.ion()
    while viewmak(shoot, framestd, vars) != 1:
        
        
        frame += 1
        plt.axis('off')
        plt.pause(0.1)
        #a=input()
        plt.cla()
        framestd = '%05d'%frame

    '''
    #plt.axis('off')
    #viewmak(shoot, framestd, vars)
    gen_png(shoot, framestd, vars)
    view_label(shoot,framestd,1)
    plt.clf()

def main(shoot, frame):
    print('{:#^10}'.format('working on shoot '+shoot+' frame '+frame))
    try:
        get_efit(shoot, frame)
    except:
        print('{:*^60}'.format('ERROR'))
        print('error shoot '+ shoot + ' frame ' + frame)
        exstr = traceback.format_exc()
        print(exstr)
        print('{:*^60}'.format(' '))


if __name__ == '__main__':
    shoots = os.listdir('efit')
    #print(shoots)
    
    start = time.time()
    with concurrent.futures.ProcessPoolExecutor(max_workers=5) as executor:
        for shoot in shoots:
            frames = os.listdir('efit\\'+shoot)
            for frame in frames:
                frame = frame.split('.mat')[0]
                #print(shoot+'    '+frame)
                executor.submit(main, shoot, frame)
    
    end = time.time()
    print('TIME: ', '%.2f'%(end - start), '   s')
    